"""
基础函数：tif文件的读写
"""

from osgeo import gdal
import numpy as np
import pandas as pd
from tqdm import tqdm
import os

DateList = pd.date_range('{}-{}'.format(2018, 12), '{}-{}'.format(2022, 11), freq='MS').strftime("%Y%m").tolist()

Geo = (113.69239932567805, 8.983152841195215e-05, 0.0, 31.36225269376597, 0.0, -8.983152841195215e-05)
Pro = """
GEOGCS["WGS 84",DATUM["WGS_1984",SPHEROID["WGS 84",6378137,298.257223563,AUTHORITY["EPSG","7030"]],AUTHORITY["EPSG","6326"]],PRIMEM["Greenwich",0,AUTHORITY["EPSG","8901"]],UNIT["degree",0.0174532925199433,AUTHORITY["EPSG","9122"]],AXIS["Latitude",NORTH],AXIS["Longitude",EAST],AUTHORITY["EPSG","4326"]]
"""

class EarlyStopping():
    """
    Early stopping to stop the training when the loss does not improve after
    certain epochs.
    """

    def __init__(self, patience=15, min_delta=0):
        """
        :param patience: how many epochs to wait before stopping when loss is
               not improving
        :param min_delta: minimum difference between new loss and old loss for
               new loss to be considered as an improvement
        """
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss):
        if self.best_loss == None:
            self.best_loss = val_loss
        elif self.best_loss - val_loss > self.min_delta:
            self.best_loss = val_loss
            # reset counter if validation loss improves
            self.counter = 0
        elif self.best_loss - val_loss < self.min_delta:
            self.counter += 1
            print(f"INFO: Early stopping counter {self.counter} of {self.patience}")
            if self.counter >= self.patience:
                print('INFO: Early stopping')
                self.early_stop = True

# 读取tiff文件
def readGeoTIFF(fileName):
    dataset = gdal.Open(fileName)
    if dataset == None:
        print(fileName + "文件无法打开")

    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    im_data = dataset.ReadAsArray(0, 0, im_width, im_height)  # 获取数据

    im_geotrans = dataset.GetGeoTransform()  # 获取仿射矩阵信息
    im_proj = dataset.GetProjection()  # 获取投影信息
    return im_data, im_geotrans, im_proj


# 写入tiff文件
def CreateGeoTiff(outRaster, image, geo_transform, projection):
    no_bands = 0
    rows = 0
    cols = 0

    driver = gdal.GetDriverByName('GTiff')
    if len(image.shape) == 2:
        no_bands = 1
        rows, cols = image.shape
    elif len(image.shape) == 3:
        no_bands, rows, cols = image.shape

    # DataSet = driver.Create(outRaster, cols, rows, no_bands, gdal.GDT_Byte)
    DataSet = driver.Create(outRaster, cols, rows, no_bands, gdal.GDT_Float32)
    DataSet.SetGeoTransform(geo_transform)
    DataSet.SetProjection(projection)

    if no_bands == 1:
        DataSet.GetRasterBand(1).WriteArray(image)  # 写入数组数据
    else:
        for i in range(no_bands):
            DataSet.GetRasterBand(i + 1).WriteArray(image[i])
    del DataSet


def Median_filtering(image, window_size=3):  # image为传入灰度图像，window_size为滤波窗口大小
    high, wide = image.shape
    img = image.copy()
    mid = (window_size - 1) // 2
    med_arry = []
    for i in tqdm(range(high - window_size)):
        for j in range(wide - window_size):
            for m1 in range(window_size):
                for m2 in range(window_size):
                    med_arry.append(int(image[i + m1, j + m2]))

            # for n in range(len(med_arry)-1,-1,-1):
            med_arry.sort()  # 对窗口像素点排序
            # print(med_arry)
            img[i + mid, j + mid] = med_arry[(len(med_arry) + 1) // 2]  # 将滤波窗口的中值赋给滤波窗口中间的像素点
            del med_arry[:]
    return img


def block_fn(x, center_val):
    unique_elements, counts_elements = np.unique(x.ravel(), return_counts=True)

    if np.isnan(center_val):
        return np.nan
    elif center_val == 1:
        return 1.0
    else:
        return unique_elements[np.argmax(counts_elements)]


def Majority_filter(x, window_size=3, type='spatial'):
    # Odd block sizes only  ( ? )
    if type == 'spatial':
        block_size = (window_size, window_size)
    else:
        block_size = (1, window_size)
    assert (block_size[0] % 2 != 0 and block_size[1] % 2 != 0)

    yy = int((block_size[0] - 1) / 2)
    xx = int((block_size[1] - 1) / 2)

    output = np.zeros_like(x)
    for i in range(0, x.shape[0]):
        miny, maxy = max(0, i - yy), min(x.shape[0] - 1, i + yy)

        for j in range(0, x.shape[1]):
            minx, maxx = max(0, j - xx), min(x.shape[1] - 1, j + xx)

            # Extract block to take majority filter over
            block = x[miny:maxy + 1, minx:maxx + 1]

            output[i, j] = block_fn(block, center_val=x[i, j])

    return output


def FilteringSpatial(img, method='NoFilter', window_size=3):
    """空间滤波"""
    if method == 'NoFilter':
        return img
    elif method == 'Median':
        return Median_filtering(img, window_size)
    elif method == 'Majority':
        return Majority_filter(img, window_size, type='spatial')


def DetectChangepoinst(x):
    """过滤变化点"""
    xx = x[0]
    id = np.where((xx[1:] - xx[:-1]) != 0)

    return id[0], np.append(xx[id], xx[-1])


def FilteringSeries(data, method='NoFilter', window_size=3):
    """时间序列滤波"""
    series = data[None, :]
    if method == 'NoFilter':
        changepoints, changetypes = DetectChangepoinst(series)
        return data, changepoints, changetypes
    elif method == 'Majority':
        res = Majority_filter(series, window_size, type='series')
        changepoints, changetypes = DetectChangepoinst(res)
        return res, changepoints, changetypes


def mat2rgb(mat):
    rgblist = [[0, 191, 255], [0, 100, 0], [60, 179, 113], [218, 165, 32], [255, 0, 0], [255, 222, 173],
               [250, 250, 210]]
    # rgblist = [[255, 255, 255], [255, 0, 0]]
    mat = mat.astype('int8')
    return np.array(rgblist)[mat]


def CreateDir(path):
    if not os.path.exists(path):
        os.makedirs(path)


if __name__ == '__main__':
    test = np.array([[0, 2, 1], [0, 2, 0], [0, 0, 0]])

    print(Median_filtering(test))
    print(Majority_filter(test))
